#include "core/main.h"

TcFeatureTransform::TcFeatureTransform(const TcFeatureTransformParams &params)
    : _params(params)
{
    numAbsPosHorzBins = params.frameWidth / params.absPosGridSize;
    numAbsPosVertBins = params.frameHeight / params.absPosGridSize;
    numAbsPosBins = numAbsPosHorzBins * numAbsPosVertBins;

    numRelPosHorzBins = 2 * (params.frameWidth / params.relPosGridSize);
    numRelPosVertBins = 2 * (params.frameHeight / params.relPosGridSize);
    numRelPosBins = numRelPosVertBins * numRelPosHorzBins;

    numSpritePairs = params.numSprites * (params.numSprites - 1);
    numSpritePairs /= 2;
    numSpritePairs += params.numSprites;

    actionBlockSize = numAbsPosBins * params.numSprites;
    actionBlockSize += numSpritePairs * numRelPosBins;
}

int TcFeatureTransform::getActionBlockStart(int actionIndex) const
{
    return actionIndex * actionBlockSize;
}

int TcFeatureTransform::getAbsPosBlockStart(int spriteId) const
{
    return spriteId * numAbsPosBins;
}

int TcFeatureTransform::getAbsPosBinIndex(const Vec2i& coord) const
{
    int bx = coord.x / _params.absPosGridSize;
    int by = coord.y / _params.absPosGridSize;
    return by * numAbsPosHorzBins + bx;
}

int TcFeatureTransform::getAbsPosBinIndex(const SpriteLocation &sprite) const
{
    return getAbsPosBlockStart(sprite.spriteID) + getAbsPosBinIndex(sprite.coord);
}

int TcFeatureTransform::getSpritePairIndex(int id1, int id2) const
{
    int higher = id1 > id2 ? id1 : id2;
    int lower = id1 < id2 ? id1 : id2;
    int index = 0;
    // This can be made constant time if it's too slow like this
    for (int i = 1; i <= higher; i++) index += i;
    index += lower;
    return index;
}

int TcFeatureTransform::getRelPosBlockStart(int id1, int id2) const
{
    int pairIndex = getSpritePairIndex(id1, id2);
    return numAbsPosBins * _params.numSprites + pairIndex * numRelPosBins;
}

int TcFeatureTransform::getRelPosBinIndex(const Vec2i &coord1, const Vec2i &coord2) const
{
    int dx = coord2.x - coord1.x;
    int dy = coord2.y - coord1.y;
    int bx = dx / _params.relPosGridSize + numRelPosHorzBins / 2;
    int by = dy / _params.relPosGridSize + numRelPosVertBins / 2;
    return by * numRelPosHorzBins + bx;
}

int TcFeatureTransform::getRelPosBinIndex(const SpriteLocation &sprite1, const SpriteLocation &sprite2) const
{
    int blockStart = getRelPosBlockStart(sprite1.spriteID, sprite2.spriteID);
    int binIndex = 0;
    if (sprite1.spriteID > sprite2.spriteID) {
        binIndex = getRelPosBinIndex(sprite1.coord, sprite2.coord);
    } else {
        binIndex = getRelPosBinIndex(sprite2.coord, sprite1.coord);
    }
    return blockStart + binIndex;
}

string TcFeatureTransform::describeFeature(int index) const
{
    stringstream ss;
    int actionIndex = index / actionBlockSize;
    ss << "action " << actionIndex << endl;

    index %= actionBlockSize;
    if (index < _params.numSprites * numAbsPosBins) {
        int spriteId = index / numAbsPosBins;
        ss << "absolute position of sprite " << spriteId << endl;
        
        index %= numAbsPosBins;
        int bx = index % numAbsPosHorzBins;
        int by = index / numAbsPosHorzBins;
        ss << "discrete position (x, y) = (" << bx << ", " << by << ")";
    } else {
        index -= _params.numSprites * numAbsPosBins;
        int pairIndex = index / numRelPosBins;
        for (int id1 = 0; id1 < _params.numSprites; id1++) {
            for (int id2 = 0; id2 <= id1; id2++) {
                if (pairIndex == getSpritePairIndex(id1, id2)) {
                    ss << "relative position of sprites " << id1 << " and " << id2 << endl;
                
                    index %= numRelPosBins;
                    int bx = index % numRelPosHorzBins;
                    int by = index / numRelPosHorzBins;

                    bx -= numRelPosHorzBins / 2;
                    by -= numRelPosVertBins / 2;

                    ss << "discrete position (dx, dy) = (" << bx << ", " << by << ")";

                    break;
                }
            }
        }
    }

    return ss.str();
}

int TcFeatureTransform::dimension() const
{
    return _params.numActions * actionBlockSize;
}

void TcFeatureTransform::transform(const Frame &frame, int actionIndex, AbstractIntVector &out) const
{
    out.clear();

    int actionBlockStart = getActionBlockStart(actionIndex);
    for (UINT i = 0; i < frame.sprites.size(); i++) {
        const SpriteLocation &sprite1 = frame.sprites[i];
        if ((unsigned long)sprite1.spriteID < (unsigned long)_params.numSprites) {
            int absPosBinIndex = getAbsPosBinIndex(sprite1) + actionBlockStart;
            out.increment(absPosBinIndex);
            for (UINT j = 0; j < i; j++) {
                const SpriteLocation &sprite2 = frame.sprites[j];
                if ((unsigned long)sprite2.spriteID < (unsigned long)_params.numSprites) {
                    int relPosBinIndex = getRelPosBinIndex(sprite1, sprite2) + actionBlockStart;
                    out.increment(relPosBinIndex);
                }
            }
        }
    }
}

void TcFeatureTransform::transform(const Frame &frame, vector<double> &out) const
{
    out.clear();
    // For now we will just encode absolute positions
    addAbsPosFeatures(frame, out);
}

// Add features encoding absolute sprite positions to the feature vector.
void TcFeatureTransform::addAbsPosFeatures(const Frame &frame, vector<double> &out) const
{
    int numAbsPosHorzBins = _params.frameWidth / _params.absPosGridSize;
    int numAbsPosVertBins = _params.frameHeight / _params.absPosGridSize;
    int absSpriteBlockSize = numAbsPosHorzBins * numAbsPosVertBins;

    UINT startSize = out.size();
    out.resize(startSize + absSpriteBlockSize * _params.numSprites, 0);

    for (auto it = frame.sprites.begin(); it != frame.sprites.end(); ++it) {
        if (it->spriteID < (UINT)_params.numSprites) {
            int bx = it->coord.x / _params.absPosGridSize;
            int by = it->coord.y / _params.absPosGridSize;

            int index = absSpriteBlockSize * it->spriteID + by * numAbsPosHorzBins + bx;
            out[startSize + index]++;
        }
    }
}

// Add features corresponding to the controller state to the feature vector.
void TcFeatureTransform::addActionFeatures(const ControllerState &action, vector<double> &out) const
{
    out.push_back(action.left() ? 1 : 0);
    out.push_back(action.right() ? 1 : 0);
    out.push_back(action.up() ? 1 : 0);
    out.push_back(action.down() ? 1 : 0);
    out.push_back(action.A() ? 1 : 0);
    out.push_back(action.B() ? 1 : 0);
}

OutputDataStream & operator<< (OutputDataStream &os, const TcFeatureTransform &transform)
{
    const TcFeatureTransformParams &params = transform._params;
    os << params.absPosGridSize
       << params.relPosGridSize
       << params.numSprites
       << params.frameHeight
       << params.frameWidth;

    return os;
}

InputDataStream & operator>> (InputDataStream &is, TcFeatureTransform &transform)
{
    return is;
}